{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style='background-image: url(\"../../share/images/header.svg\") ; padding: 0px ; background-size: cover ; border-radius: 5px ; height: 250px'>\n",
    "    <div style=\"float: right ; margin: 50px ; padding: 20px ; background: rgba(255 , 255 , 255 , 0.7) ; width: 50% ; height: 150px\">\n",
    "        <div style=\"position: relative ; top: 50% ; transform: translatey(-50%)\">\n",
    "            <div style=\"font-size: xx-large ; font-weight: 900 ; color: rgba(0 , 0 , 0 , 0.8) ; line-height: 100%\">Computational Seismology</div>\n",
    "            <div style=\"font-size: large ; padding-top: 20px ; color: rgba(0 , 0 , 0 , 0.5)\">Spectral Element Method - 1D Elastic Wave Equation</div>\n",
    "        </div>\n",
    "    </div>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<p style=\"width:20%;float:right;padding-left:50px\">\n",
    "<img src=../../share/images/book.jpg>\n",
    "<span style=\"font-size:smaller\">\n",
    "</span>\n",
    "</p>\n",
    "\n",
    "\n",
    "---\n",
    "\n",
    "This notebook is part of the supplementary material \n",
    "to [Computational Seismology: A Practical Introduction](https://global.oup.com/academic/product/computational-seismology-9780198717416?cc=de&lang=en&#), \n",
    "Oxford University Press, 2016.\n",
    "\n",
    "\n",
    "##### Authors:\n",
    "* David Vargas ([@dvargas](https://github.com/davofis))\n",
    "* Heiner Igel ([@heinerigel](https://github.com/heinerigel))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic Equations\n",
    "\n",
    "This notebook presents the numerical solution for the 1D elastic wave equation\n",
    "\n",
    "\\begin{equation}\n",
    "\\rho(x) \\partial_t^2 u(x,t) = \\partial_x (\\mu(x) \\partial_x u(x,t)) + f(x,t),\n",
    "\\end{equation}\n",
    "\n",
    "using the spectral element method. This is done after a series of steps summarized as follow:\n",
    "\n",
    "1) The wave equation is written into its Weak form\n",
    "\n",
    "2) Apply stress Free Boundary Condition after integration by parts \n",
    "\n",
    "3) Approximate the wave field as a linear combination of some basis\n",
    "\n",
    "\\begin{equation}\n",
    "u(x,t) \\ \\approx \\ \\overline{u}(x,t) \\ = \\ \\sum_{i=1}^{n} u_i(t) \\ \\varphi_i(x)\n",
    "\\end{equation}\n",
    "\n",
    "4) Use the same basis functions in $u(x, t)$ as test functions in the weak form, the so call Galerkin principle.\n",
    "\n",
    "6) The continuous weak form is written as a system of linear equations by considering the approximated displacement field.\n",
    "\n",
    "\\begin{equation}\n",
    "\\mathbf{M}^T\\partial_t^2 \\mathbf{u} + \\mathbf{K}^T\\mathbf{u} = \\mathbf{f}\n",
    "\\end{equation}\n",
    "\n",
    "7) Time extrapolation with centered finite differences scheme\n",
    "\n",
    "\\begin{equation}\n",
    "\\mathbf{u}(t + dt) = dt^2 (\\mathbf{M}^T)^{-1}[\\mathbf{f} - \\mathbf{K}^T\\mathbf{u}] + 2\\mathbf{u} - \\mathbf{u}(t-dt).\n",
    "\\end{equation}\n",
    "\n",
    "where $\\mathbf{M}$ is known as the mass matrix, and $\\mathbf{K}$ the stiffness matrix.\n",
    "\n",
    "The above solution is exactly the same presented for the classic finite-element method. Now we introduce appropriated basis functions and integration scheme to efficiently solve the system of matrices.\n",
    "\n",
    "#### Interpolation with Lagrange Polynomials\n",
    "At the elemental level (see section 7.4), we introduce as interpolating functions the Lagrange polynomials and use $\\xi$ as the space variable representing our elemental domain:\n",
    "\n",
    "\\begin{equation}\n",
    "\\varphi_i \\ \\rightarrow \\ \\ell_i^{(N)} (\\xi) \\ := \\ \\prod_{j \\neq i}^{N+1} \\frac{\\xi - \\xi_j}{\\xi_i-\\xi_j}, \\qquad   i,j = 1, 2, \\dotsc , N + 1  \n",
    "\\end{equation}\n",
    "\n",
    "#### Numerical Integration\n",
    "The integral of a continuous function $f(x)$ can be calculated after replacing $f(x)$ by a polynomial approximation that can be integrated analytically. As interpolating functions we use again the Lagrange polynomials and\n",
    "obtain Gauss-Lobatto-Legendre quadrature. Here, the GLL points are used to perform the integral. \n",
    "\n",
    "\\begin{equation}\n",
    "\\int_{-1}^1 f(x) \\ dx \\approx \\int _{-1}^1 P_N(x) dx = \\sum_{i=1}^{N+1}\n",
    " w_i f(x_i) \n",
    "\\end{equation}\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0
    ],
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Import all necessary libraries, this is a configuration step for the exercise.\n",
    "# Please run it before the simulation code!\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "# Show Plot in The Notebook\n",
    "matplotlib.use(\"nbagg\")\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from gll import gll\n",
    "from lagrange1st import lagrange1st \n",
    "from ricker import ricker\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Initialization of setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Initialization of setup\n",
    "# ---------------------------------------------------------------\n",
    "nt    = 10000         # number of time steps\n",
    "xmax  = 10000.        # Length of domain [m]\n",
    "vs    = 2500.         # S velocity [m/s]\n",
    "rho   = 2000          # Density [kg/m^3]\n",
    "mu    = rho * vs**2   # Shear modulus mu\n",
    "N     = 3             # Order of Lagrange polynomials\n",
    "ne    = 250           # Number of elements\n",
    "Tdom  = .2            # Dominant period of Ricker source wavelet\n",
    "iplot = 20            # Plotting each iplot snapshot\n",
    "\n",
    "# variables for elemental matrices\n",
    "Me = np.zeros(N+1, dtype =  float)\n",
    "Ke = np.zeros((N+1, N+1), dtype =  float)\n",
    "# ----------------------------------------------------------------\n",
    "\n",
    "# Initialization of GLL points integration weights\n",
    "[xi, w] = gll(N)    # xi, N+1 coordinates [-1 1] of GLL points\n",
    "                    # w Integration weights at GLL locations\n",
    "# Space domain\n",
    "le = xmax/ne        # Length of elements\n",
    "# Vector with GLL points  \n",
    "k = 0\n",
    "xg = np.zeros((N*ne)+1) \n",
    "xg[k] = 0\n",
    "for i in range(1,ne+1):\n",
    "    for j in range(0,N):\n",
    "        k = k+1\n",
    "        xg[k] = (i-1)*le + .5*(xi[j+1]+1)*le\n",
    "\n",
    "# ---------------------------------------------------------------\n",
    "dxmin = min(np.diff(xg))  \n",
    "eps = 0.1           # Courant value\n",
    "dt = eps*dxmin/vs   # Global time step\n",
    "\n",
    "# Mapping - Jacobian\n",
    "J = le/2 \n",
    "Ji = 1/J    # Inverse Jacobian\n",
    "\n",
    "# 1st derivative of Lagrange polynomials\n",
    "l1d = lagrange1st(N)   # Array with GLL as columns for each N+1 polynomial\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. The Mass Matrix\n",
    "Now we initialize the mass and stiffness matrices. In general, the mass matrix at the elemental level is given\n",
    "\n",
    "\\begin{equation}\n",
    "M_{ji}^e \\ = \\ w_j  \\ \\rho (\\xi)  \\ \\frac{\\mathrm{d}x}{\\mathrm{d}\\xi} \\delta_{ij} \\vert_ {\\xi = \\xi_j}\n",
    "\\end{equation}\n",
    "\n",
    "#### Exercise 1 \n",
    "Implement the mass matrix using the integration weights at GLL locations $w$, the jacobian $J$, and density $\\rho$. Then, perform the global assembly of the mass matrix, compute its inverse, and display the inverse mass matrix to visually inspect how it looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "code_folding": [],
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#################################################################\n",
    "# IMPLEMENT THE MASS MATRIX HERE!\n",
    "#################################################################\n",
    "\n",
    "\n",
    "#################################################################\n",
    "# PERFORM THE GLOBAL ASSEMBLY OF M HERE!\n",
    "#################################################################\n",
    "\n",
    "\n",
    "#################################################################\n",
    "# COMPUTE THE INVERSE MASS MATRIX HERE!\n",
    "#################################################################\n",
    "\n",
    "\n",
    "#################################################################\n",
    "# DISPLAY THE INVERSE MASS MATRIX HERE!\n",
    "#################################################################\n",
    "    \n",
    "     "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. The Stiffness matrix\n",
    "On the other hand, the general form of the stiffness matrix at the elemtal level is\n",
    "\n",
    "\\begin{equation}\n",
    "K_{ji}^e \\ = \\ \\sum_{k = 1}^{N+1} w_k \\mu (\\xi) \\partial_\\xi \\ell_j (\\xi) \\partial_\\xi \\ell_i (\\xi) \\left(\\frac{\\mathrm{d}\\xi}{\\mathrm{d}x} \\right)^2 \\frac{\\mathrm{d}x}{\\mathrm{d}\\xi} \\vert_{\\xi = \\xi_k}\n",
    "\\end{equation} \n",
    "\n",
    "#### Exercise 2 \n",
    "Implement the stiffness matrix using the integration weights at GLL locations $w$, the jacobian $J$, and shear stress $\\mu$. Then, perform the global assembly of the mass matrix and display the matrix to visually inspect how it looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#################################################################\n",
    "# IMPLEMENT THE STIFFNESS MATRIX HERE!\n",
    "#################################################################\n",
    "\n",
    "\n",
    "#################################################################\n",
    "# PERFORM THE GLOBAL ASSEMBLY OF K HERE!\n",
    "#################################################################\n",
    "\n",
    "    \n",
    "#################################################################\n",
    "# DISPLAY THE STIFFNESS MATRIX HERE!\n",
    "#################################################################\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Finite element solution \n",
    "\n",
    "Finally we implement the spectral element solution using the computed mass $M$ and stiffness $K$ matrices together with a finite differences extrapolation scheme\n",
    "\n",
    "\\begin{equation}\n",
    "\\mathbf{u}(t + dt) = dt^2 (\\mathbf{M}^T)^{-1}[\\mathbf{f} - \\mathbf{K}^T\\mathbf{u}] + 2\\mathbf{u} - \\mathbf{u}(t-dt).\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     39
    ],
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# SE Solution, Time extrapolation\n",
    "# ---------------------------------------------------------------\n",
    "\n",
    "# initialize source time function and force vector f\n",
    "src  = ricker(dt,Tdom)\n",
    "isrc = int(np.floor(ng/2))   # Source location\n",
    "\n",
    "# Initialization of solution vectors\n",
    "u = np.zeros(ng)\n",
    "uold = u\n",
    "unew = u\n",
    "f = u \n",
    "\n",
    "# Initialize animated plot\n",
    "# ---------------------------------------------------------------  \n",
    "plt.figure(figsize=(10,6))\n",
    "lines = plt.plot(xg, u, lw=1.5)\n",
    "plt.title('SEM 1D Animation', size=16)\n",
    "plt.xlabel(' x (m)')\n",
    "plt.ylabel(' Amplitude ')\n",
    "\n",
    "plt.ion() # set interective mode\n",
    "plt.show()\n",
    "\n",
    "# ---------------------------------------------------------------\n",
    "# Time extrapolation\n",
    "# ---------------------------------------------------------------\n",
    "for it in range(nt): \n",
    "    # Source initialization\n",
    "    f= np.zeros(ng)\n",
    "    if it < len(src):\n",
    "        f[isrc-1] = src[it-1] \n",
    "               \n",
    "    # Time extrapolation\n",
    "    unew = dt**2 * Minv @ (f - K @ u) + 2 * u - uold\n",
    "    uold, u = u, unew\n",
    "\n",
    "    # --------------------------------------   \n",
    "    # Animation plot. Display solution \n",
    "    if not it % iplot:\n",
    "        for l in lines:\n",
    "            l.remove()\n",
    "            del l\n",
    "        # -------------------------------------- \n",
    "        # Display lines            \n",
    "        lines = plt.plot(xg, u, color=\"black\", lw = 1.5)\n",
    "        plt.gcf().canvas.draw()       "
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
